package utils

import (
	"sync/atomic"
	"time"
)

/*
**

	TryBegin() 如果失败直接返回, 该方式比Chan效率高 (1000 抢 10个计数)
	Begin()    会阻塞比较, 该方式比Chan效率高 (1000 抢10个计数) 缺点 经常到不了10个 纤程
	  缺陷: Count() 可能会超过MAX大小
*/
type BlockEx struct {
	cnt int32
	max int32
}

func NewBlockEx(max int32) *BlockEx {
	rval := &BlockEx{
		cnt: 0,
		max: max,
	}
	return rval
}

func (this *BlockEx) TryBegin() bool {
	if this.max == 0 {
		atomic.AddInt32(&this.cnt, 1)
		return true
	}

	r := this.cnt
	if r >= this.max {
		return false
	}

	r = atomic.AddInt32(&this.cnt, 1)
	if r <= this.max {
		return true
	} else {
		atomic.AddInt32(&this.cnt, -1)
		return false
	}
}

func (this *BlockEx) Begin() {
	for {
		if this.TryBegin() {
			break
		} else {
			time.Sleep(time.Millisecond)
		}
	}
}

/*
*

	在大并发情况下, cnt可能会超过 max, 原因是在tryBegin时, 会尝试+1, 如果计数大于 max才会减一 造成了cnt读取那一瞬间会造成多的情况
*/
func (this *BlockEx) Count() int {
	return int(this.cnt)
}

func (this *BlockEx) End() {
	atomic.AddInt32(&this.cnt, -1)
}
